import numpy as np
import faiss

class FaissKNeighbors:
    def __init__(self, k=1, res=None):
        self.index = None
        self.y = None
        self.k = k
        self.res = res

    def fit(self, X, y):
        # 初始化 self.index 为一个FAISS索引: IndexFlatL2, 该索引使用欧氏距离进行搜索
        # 如果有GPU资源对象，则将索引转移到GPU上
        self.index = faiss.IndexFlatL2(X.shape[1])
        if self.res is not None:
            self.index = faiss.index_cpu_to_gpu(self.res, 0, self.index)
        
        # 将训练数据添加到索引中
        self.index.add(X.astype(np.float32))
        
        # 初始化 self.y 为传入的 y
        self.y = y

    def predict(self, X):
        distances, indices = self.index.search(X.astype(np.float32), self.k)
        votes = self.y[indices]
        predictions = np.array([np.argmax(np.bincount(vote)) for vote in votes])
        return predictions

    def score(self, X, y):
        # 预测并比较预测结果和真实标签，计算准确率
        predictions = self.predict(X)
        accuracy = np.mean(predictions == y)
        return accuracy